import segyio


# 从指定的起始索引开始，查找第一个非零值的时间索引
def find_first_non_zero_time(data, start_index):
    """
    从指定的起始索引开始，查找第一个非零值的时间索引。
    :param data: 数据数组
    :param start_index: 起始索引
    :return: 第一个非零值的时间索引，如果没有非零值，返回None
    """
    for i in range(start_index, len(data)):
        if data[i] != 0:
            return i
    return None


# 判断sgy文件的数据是向左偏移还是向右偏移
def get_flag(segy_file, start_trace):
    """
        判断sgy文件的数据是向左偏移还是向右偏移
        :param segy_file: SEG-Y文件数据
        :param start_trace: 起始道编号
        :return: 偏移标志，如果是-1，则向左偏移，如果是1，则向右偏移
        """
    flag = -1  # 如果是-1，则向左偏移，如果是1，则向右偏移

    # 从指定道开始处理
    trace_index = start_trace - 1  # segyio的道索引从0开始
    trace_data = segy_file.trace[trace_index]  # 读取当前道的数据
    start_time_index = 0  # 第1道开始时间为0
    non_zero_time_index = find_first_non_zero_time(trace_data, start_time_index)  # 查找第一个非零值的时间索引
    # 左侧第2道数据
    trace_index2 = trace_index + flag
    trace_data2 = segy_file.trace[trace_index2]  # 读取当前道的数据
    start_time_index2 = non_zero_time_index  # 开始时间为上一道第一个不为0的时间索引
    non_zero_time_index2 = find_first_non_zero_time(trace_data2, start_time_index2)  # 查找第一个非零值的时间索引
    # 判断左偏还是右偏
    if non_zero_time_index == non_zero_time_index2:
        flag = 1

    return flag


# 处理SEG-Y文件，从指定道开始，逐道查找第一个非零值的时间
def get_trace_offset(segy_file, start_trace, trace_num):
    """
    处理SEG-Y文件，从指定道开始，逐道查找第一个非零值的时间。
    :param segy_file: SEG-Y文件数据
    :param start_trace: 起始道编号
    :param trace_num: 走廊通道数
    :return: 道偏移的相关数据，
    [
        {
        "trace_index"，道编号，从1开始
        "non_zero_time"，该道第一个非零点的时间
        },
        ...
    ]
    """
    sample_interval = segy_file.bin[segyio.BinField.Interval] / 1000  # 获取采样间隔（ms），转换为ms
    flag = get_flag(segy_file, start_trace)  # 获取偏移方向
    trace_offset = []  # 存储道偏移的相关数据

    if flag == 1:  # 向右偏移时，开始计算的初始道为最右侧的道
        start_trace = start_trace + trace_num + 1

    # 从指定道开始处理
    trace_index = start_trace + flag  # segyio的道索引从0开始
    last_non_zero_time = 0

    while trace_index >= 0:
        trace_data = segy_file.trace[trace_index]  # 读取当前道的数据
        if trace_index == start_trace - 1:  # 如果是第一道，从0ms开始排查
            start_time_index = 0
        else:  # 否则从上一道记录的时间开始排查
            start_time_index = last_non_zero_time

        # 查找第一个非零值的时间索引
        non_zero_time_index = find_first_non_zero_time(trace_data, start_time_index)
        if non_zero_time_index is not None:  # 有非零值
            non_zero_time = non_zero_time_index * sample_interval  # 通过索引和采样间隔得到事件
            print(f"道 {trace_index + 1}: 第一个非零值的时间为 {non_zero_time} ms")
            last_non_zero_time = non_zero_time_index
            trace_offset.append({
                "trace_index": trace_index + 1,  # 道编号
                "non_zero_time": non_zero_time  # 该道第一个非零点的时间
            })
        else:  # 没有找到非零值
            print(f"道 {trace_index + 1}: 没有找到非零值")
            break

        # 准备处理下一道
        trace_index = trace_index + flag

    return trace_offset


# 通过裁剪得到走廊数据，包括start_index和end_index
def clip_trace(inp_file, out_file, start_index, end_index):
    """
    通过裁剪得到走廊数据，包括start_index和end_index
    :param inp_file: 被截取的sgy文件路径，地面镶嵌数据
    :param out_file: 截取后保存的sgy文件路径，走廊
    :param start_index: 截取道的起始下标，地面镶嵌数据中走廊的起始位置
    :param end_index: 截取道的结束下标，地面镶嵌数据中走廊的结束位置
    """
    with segyio.open(inp_file, 'r', ignore_geometry=True) as infile:
        # 以写入模式打开输出的 SGY 文件
        spec = segyio.spec()
        spec.sorting = infile.sorting
        spec.format = infile.format
        spec.samples = infile.samples
        spec.tracecount = end_index - start_index + 1  # 数据总道数
        with segyio.create(out_file, spec) as outfile:
            # 复制二进制头信息
            outfile.bin = infile.bin
            # 复制道头信息和数据
            for i, trace_index in enumerate(range(start_index - 1, end_index)):
                outfile.header[i] = infile.header[trace_index]
                outfile.trace[i] = infile.trace[trace_index]


def get_trace_data(file_path, trace_offset, flag):
    with segyio.open(file_path, strict=False, ignore_geometry=True) as segy_file:
        sample_interval = segy_file.bin[segyio.BinField.Interval] / 1000  # 获取采样间隔（ms），转换为ms
        if flag == -1:
            # 左侧数据
            trace_data = []
            trace_offset_left = trace_offset[::-1]
            for i in range(segy_file.tracecount):
                if i < len(trace_offset_left):
                    time_index = int(trace_offset[i]["non_zero_time"] / sample_interval)
                    data = segy_file.trace[i][time_index:]
                else:
                    data = segy_file.trace[i]
                trace_data.append(data)

    return 0

# 将走廊、滤波走廊、地面数据从镶嵌数据中裁剪出来
if __name__ == "__main__":
    file_path = r"E:\vspData\斜井\psdm_stk_inl994_1394_0914vsp_mu_filter_scale_inq_3D_tm_inlne1194AGC镶嵌1.sgy"
    stkbf_path = r"E:\vspData\斜井\01滤波走廊.sgy"
    stk_path = r"E:\vspData\斜井\01走廊.sgy"
    sein_path = r"E:\vspData\斜井\01地面.sgy"
    trace_num = 10
    sein_num = 50
    trace_start = 404

    with segyio.open(file_path, strict=False, ignore_geometry=True) as segy_file:
        trace_offset = get_trace_offset(segy_file, trace_start, trace_num)  # 道偏移数据
        flag = get_flag(segy_file, trace_start)  # 数据向左偏还是向右偏

        # 截取滤波走廊
        if flag == -1:  # 向左偏移
            start_index = trace_offset[len(trace_offset) - 1]["trace_index"]  # 截取道的开始下标，392
            end_index = trace_start + trace_num  # 截取道的结束下标，414
        else:  # 向右偏移
            start_index = trace_start + 1  # 截取道的开始下标，405
            end_index = trace_offset[len(trace_offset) - 1]["trace_index"]  # 截取道的结束下标,427
        clip_trace(file_path, stkbf_path, start_index, end_index)  # 裁剪得到滤波走廊数据

        # 截取走廊
        start_index2 = start_index + trace_num + 1  # 截取道的开始下标，403，416
        end_index2 = end_index + trace_num + 1  # 截取道的结束下标，425，438
        clip_trace(file_path, stk_path, start_index2, end_index2)  # 裁剪得到滤波走廊数据

        # 截取地面
        start_index3 = start_index2 + trace_num + 1  # 截取道的开始下标，414，427
        end_index3 = end_index2 + sein_num + 1  # 截取道的结束下标，476，489
        clip_trace(file_path, sein_path, start_index3, end_index3)  # 裁剪得到滤波走廊数据

        # 将地面嵌入数据的相关偏移数据，改为裁剪出来的走廊的相关偏移数据，trace_index从1开始
        for i in range(len(trace_offset)):
            trace_offset[i]["trace_index"] = trace_offset[i]["trace_index"] - start_index + 1

        stkbf_data = get_trace_data(stkbf_path, trace_offset, flag)